import numpy as np
from numpy import dot
from numpy.linalg import norm


def euclidean(x,y):
    '''欧式距离'''
    d = np.asarray(x) - np.asarray(y)
    return norm(d)

def cosine(x,y):
    '''余弦距离'''
    d = dot(x,y) / (norm(x) * norm(y))
    return d


import torch
from transformers import AutoModel, AutoTokenizer

# 初始化
model = AutoModel.from_pretrained(r"D:\HuaweiMoveData\Users\86157\Desktop\WitNova\WitNova_backend\m3e-base", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(r"D:\HuaweiMoveData\Users\86157\Desktop\WitNova\WitNova_backend\m3e-base", trust_remote_code=True)

# 输入文本
# text = "自然语言处理技术"

# 生成向量函数
def get_embedding(text):
    inputs = tokenizer(
        text,
        return_tensors="pt",  # 返回PyTorch张量
        padding = True,  # 自动填充到最长序列
        truncation = True,  # 自动截断到模型最大长度
        max_length=512  # 显式设置最大长度
    )

    # 修改后的分批处理
    batch_size = 8  # 根据内存调整
    embeddings = []
    for i in range(0, len(inputs.input_ids), batch_size):
        batch = {
            "input_ids": inputs.input_ids[i:i + batch_size],
            "attention_mask": inputs.attention_mask[i:i + batch_size]
        }
        with torch.no_grad():
            outputs = model(**batch)
        embeddings.append(outputs.last_hidden_state[:, 0])
    embeddings = torch.cat(embeddings)

    # return outputs.last_hidden_state.mean(dim=1).cpu().numpy()
    return embeddings.detach().cpu().numpy()  # 转换为numpy数组

# vec = get_embedding(text)
# print(vec)


# from sklearn.metrics.pairwise import cosine_similarity
#
# text1 = "这个多少钱"
# text2 = "这个什么价格"
# text3 = "这个给我吧"
#
# vec1 = get_embedding(text1)
# vec2 = get_embedding(text2)
# vec3 = get_embedding(text3)
#
# similarity1 = cosine_similarity(vec1, vec2)[0][0]
# print(f"1&2相似度: {similarity1:.4f}")  # 示例输出: 0.7823
# similarity2 = cosine_similarity(vec1, vec3)[0][0]
# print(f"1&3相似度: {similarity2:.4f}")
# similarity3 = cosine_similarity(vec2, vec3)[0][0]
# print(f"2&3相似度: {similarity3:.4f}")


# similarity1 = euclidean(vec1, vec2)
# print(f"1&2欧氏距离: {similarity1:.4f}")
# similarity2 = euclidean(vec1, vec3)
# print(f"1&3欧氏距离: {similarity2:.4f}")
# similarity3 = euclidean(vec2, vec3)
# print(f"2&3欧氏距离: {similarity3:.4f}")